"""
https://www.tensorflow.org/guide/keras/train_and_evaluate
"""

import tensorflow as tf
import tensorflow_addons as tfa
from tensorflow import keras
from tensorflow.keras import layers, losses, metrics, optimizers, Model
from python_ai.common.xcommon import *
import os
import sys
from sklearn.datasets import load_boston
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler

# data
x, y = load_boston(return_X_y=True)
scaler = StandardScaler()
x = scaler.fit_transform(x)
y = scaler.fit_transform(y.reshape([-1, 1]))

x_train, x_test, y_train, y_test = train_test_split(x, y, train_size=0.7, random_state=777)
m, n = x_train.shape
sep((m, n))

# Reserve 10,000 samples for validation
n_val = int(np.ceil(m * 0.1))
x_val = x_train[-n_val:]
y_val = y_train[-n_val:]
x_train = x_train[:-n_val]
y_train = y_train[:-n_val]

# model
L1 = 400
L2 = 200
n_epoch = 10
batch_size = 128
inputs = keras.Input(shape=(n,), name="inputs")
x = layers.Dense(L1, activation="relu", name="dense_1")(inputs)
x = layers.Dense(L2, activation="relu", name="dense_2")(x)
outputs = layers.Dense(1, activation=None, name="predictions")(x)  # ATTENTION not activation when do lin regr rather than logistic regr

model = Model(inputs=inputs, outputs=outputs)

# specify the training config


class MyR2(metrics.Metric):

    def __init__(self, name='r2', **kwargs):
        super(MyR2, self).__init__(name=name, **kwargs)
        self.u = tf.Variable(0., dtype=tf.float32, name='n_total')
        self.v = tf.Variable(0., dtype=tf.float32, name='n_total')

    def reset_state(self):
        self.u.assign(0.)
        self.v.assign(0.)

    def result(self):
        return 1 - self.u / self.v

    def update_state(self, y_true, y_pred, sample_weight=None):
        y_m = tf.reduce_mean(y_true)
        self.u.assign_add(tf.reduce_mean((y_true - y_pred) ** 2))
        self.v.assign_add(tf.reduce_mean((y_true - y_m) ** 2))


model.compile(
    optimizer=optimizers.RMSprop(),
    loss=losses.mean_squared_error,
    metrics=[
        metrics.mean_squared_error,
        MyR2(),
        tfa.metrics.r_square.RSquare(name='r_square', y_shape=(1,)),
        metrics.mean_absolute_percentage_error,
        metrics.mean_absolute_error
    ]
)

# fit and train
sep('Fit and train')
check_shape(x_train, 'x_train')
check_shape(y_train, 'y_train')
check_shape(x_val, 'x_val')
check_shape(y_val, 'y_val')
history = model.fit(
    x_train,
    y_train,
    batch_size=batch_size,
    epochs=n_epoch,
    validation_data=(x_val, y_val)
)
print(history.history)

sep('Evaluate on test data')
results = model.evaluate(x_test, y_test, batch_size=128)
print(f'test loss, test acc: {results}')

